package org.fujay.commons.langchain4j.rag.augmentor.aggregator;

import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.langchain4j.rag.augmentor.EnhancedContentAggregator;

import java.util.*;

import static dev.langchain4j.internal.ValidationUtils.ensureBetween;

/**
 * @author slm
 * @description 基于RRF（Reciprocal Rank Fusion）算法的内容聚合器
 * 适合聚合多个检索器返回的结果，能有效处理不同分数范围的结果
 */
@Slf4j
public class RrfContentAggregator implements EnhancedContentAggregator {

    private final int k;
    private final int maxResults;

    /**
     * 创建默认参数的RRF聚合器
     * 默认k值为60，最大结果数为10
     */
    public RrfContentAggregator() {
        this(60, 10);
    }

    /**
     * 创建自定义参数的RRF聚合器
     *
     * @param k          RRF算法中的常数参数，用于平滑排名差异，通常设置为60
     *                   较大的k值会减弱排名差异，较小的k值会放大排名靠前内容的重要性
     *                   k必须大于等于1
     * @param maxResults 返回的最大结果数
     */
    public RrfContentAggregator(int k, int maxResults) {
        ensureBetween(k, 1, Integer.MAX_VALUE, "k");
        this.k = k;
        this.maxResults = maxResults;
    }

    /**
     * Aggregates all {@link Content}s retrieved by all {@link ContentRetriever}s
     * using all {@link Query}s.
     * The {@link Content}s, both on input and output, are sorted by relevance,
     * with the most relevant {@link Content}s appearing at the beginning of
     * {@code List<Content>}.
     *
     * @param queryToContents A map from a {@link Query} to all
     *                        {@code List<Content>} retrieved with that
     *                        {@link Query}.
     *                        Given that each {@link Query} can be routed to
     *                        multiple {@link ContentRetriever}s, the
     *                        value of this map is a
     *                        {@code Collection<List<Content>>}
     *                        rather than a simple {@code List<Content>}.
     * @return A list of aggregated {@link Content}s.
     */
    @Override
    public List<Content> aggregate(Map<Query, Collection<List<Content>>> queryToContents) {
        if (queryToContents == null || queryToContents.isEmpty()) {
            return Collections.emptyList();
        }

        // 汇总所有查询获取的内容
        List<List<Content>> allContentLists = new ArrayList<>();
        for (Collection<List<Content>> contentLists : queryToContents.values()) {
            allContentLists.addAll(contentLists);
        }

        // 使用RRF算法融合内容
        List<Content> fusedContents = fuse(allContentLists);

        // 限制结果数量
        return fusedContents.size() <= maxResults
                ? fusedContents
                : fusedContents.subList(0, maxResults);
    }

    /**
     * 使用RRF算法融合多个内容列表
     *
     * @param listsOfContents 要融合的多个内容列表
     * @return 融合后的内容列表
     */
    private List<Content> fuse(Collection<List<Content>> listsOfContents) {
        Map<Content, Double> scores = new LinkedHashMap<>();

        for (List<Content> contentList : listsOfContents) {
            for (int i = 0; i < contentList.size(); i++) {
                Content content = contentList.get(i);
                double currentScore = scores.getOrDefault(content, 0.0);
                int rank = i + 1;
                double newScore = currentScore + 1.0 / (k + rank);
                scores.put(content, newScore);
            }
        }

        List<Content> fused = new ArrayList<>(scores.keySet());
        fused.sort(Comparator.comparingDouble(scores::get).reversed());

        return fused;
    }
}
